namespace import Path::expand 
namespace import Path::VarSpace
namespace eval SimTask {
# possible Environment : cluster single threaded
variable Environment cluster
variable Njob  12
variable JobHeader "#!/bin/bash\n#PBS -q production\n#PBS -N NLie\n#PBS -l select=1:ncpus=1\n#PBS -l place=pack\n#PBS -V"

namespace export *

# setup 
proc check_num_cpu {} {}
proc check_user {} {}
proc check_if_pbs {} {}
proc check_queue {} {}
proc setup {} {}


# prepare structures
proc prepfiles { pdb_id files} {}
proc gen { path name seq {genscript gen.inp}} {}

# run simulations
proc run { n base script optlist } {}
# restarting simulations
proc restart { n base script optlist checkfile {test -notest} } {}
# analyzing dcd files
proc ana { {base ""} dcdfile script options  {checkrange -nocheck} } {}
# helper
proc check_queue {} {}
proc submit_or_wait { job {queue 12} {wait 60} } {}


}

#short cut
proc SimTask::gen { path name seq {genscript gen.inp}} {
    set fullseq [ expand_seq $seq ]
    set seqlength [llength $fullseq ]
    variable template_path 
    variable base
    set seqlen $seqlength
    set seq $fullseq
    VarSpace option { seqlen seq} 
    run 1 $path/$name $genscript option
}

proc SimTask::genOligomer { } {

}

proc SimTask::check_if_threaded {} {
 if [catch {set version [ package require Thread] }]  {
   return 0
 } else {
   return 1
 }
}
proc SimTask::check_user {} {
   return $::env(USER)
}
proc SimTask::check_num_cpu {} {
  if [catch { set cpulist [exec grep -i "^processor" /proc/cpuinfo]}] {
   puts 'Cannot determine number of cpus!'
   return 1
  } else {
    set has [llength [split $cpulist \n]]
    return $has
  }
}
proc SimTask::check_if_pbs {} {
  if [ catch { set queue [exec qstat -u [check_user]]}] {
   return 0
  } else {
   return 1
  }
}
proc SimTask::check_queue {} {
  if [ catch { set queue [exec qstat -u [check_user]]}] {
   return -1
  } else {
   set i 0
   foreach line [split $queue \n] {
   if {[string first [check_user] $line]> -1} { incr i}
   }
   return $i
  }
}
proc SimTask::check_cmd { alias} {
  set bins {}
  foreach name $alias {
  append bins " [check_bin $name] "
  }
  if {[llength $bins] <= 0} { return "Not found"}
  return $bins
}
proc SimTask::check_bin { cmd } {
 if [catch {exec which $cmd}] {
  return ""
 } else {
  return [exec which $cmd]
 }
}

proc SimTask::check_remainning_workpower {} {
    variable Environment 
    variable Njobs
    set ntrial 5
    switch $Environment {
	cluster {
	    set remained [expr $Njobs - [SimTask::check_queue ]]
	} 
	threaded {
	    set inqueue 0
            for { set i 0 } { $i < $ntrial } { incr i } {
	    set result [exec ps -Af r ]
	    foreach line [split $result \n] {
		set datalist [regexp -all -inline -- {\S+} $line]
                set user [lindex $datalist 0]
                set cmd [lindex $datalist 8]
		set stat [lindex $datalist 6]
		if {([string match R* $stat] == 1 || [string match D* $stat] == 1) && ![string match ps $cmd] } {
		    incr inqueue
		}
	    }
            after 50
            }
	    set remained [expr $Njobs - $inqueue/$ntrial]
	}
    }
    return $remained
}
proc SimTask::submit_or_wait {job {waittime 10} } {
    variable Environment
    variable RunningTask
    if {[check_remainning_workpower] <= -18 } { error "Computational resource critically low!!!"} 
    while { [check_remainning_workpower] <= 0 } {
        puts "Queue Full: No resource available"
	after [expr $waittime *1000]
    }
    set remaining [check_remainning_workpower] 
    puts "$remaining CPU left."
    if { $remaining < 0 } { puts "Warning: problem reading job list:\n[exec ps -Af r ]" }
    puts "Submitting job $job"
    set caller [dict get [info frame -1] proc ]
    switch $Environment {
	cluster {
	   lappend RunningTask($caller) [lindex [split [exec qsub -o $job.out -e $job.err $job] .] 0 ]
	}
	default {
	   lappend RunningTask($caller) [exec bash $job 1>$job.out 2>$job.err & ]
            after 100
	}
    } 
}

proc SimTask::waitfor { proc } {
    variable RunningTask
    set alldone false
    # if proc start with ::SimTask
    if {[string first ::SimTask $proc ] < 0} { set proc ::SimTask::$proc }
    if ![info exists RunningTask($proc)] return
    while {! $alldone } {
        puts "Wait for $proc ..."
            set alldone true
            for {set i 0} {$i < [llength $RunningTask($proc)] } {incr i} {
                 if [checkIfTaskExist [lindex $RunningTask($proc) $i]] {
                    set alldone false 
                 } 
            }
            after 5000
    }
}

#proc SimTask::holdUntilAllFinished {} {
#    variable RunningTask
#    set alldone false
#    while {! $alldone } {
#            set alldone true
#            for {set i 0} {$i < [llength $RunningTask] } {incr i} {
#                 if [checkIfTaskExist [lindex $RunningTask $i]] {
#                    set alldone false 
#                 } 
#            }
#            after 5000
#    }
#}
proc SimTask::checkIfTaskExist { id} {
    variable Environment
    variable RunningTask
    switch $Environment {
	cluster {
        if [ catch {set result [exec qstat $id ]}] {return false }
        if {[string first  "Unknown Job Id" $result] > 0} {
            return false
        } else {
            return true
        }
	}
	default {
       if [catch {set result [ exec ps -p $id ]}] {return false }
       if {[string first $id $result ] > 0} {
          return true
       } else {
          return false
       }
	}
    } 
}
proc SimTask::setup {} {
    variable Njobs
    variable Environment
    variable Charmmbin 
    variable Vmdbin
    variable Tclbin
    variable PdbPath
    variable JobPath
    variable TemplatePath
    variable RunPath

    set Environment threaded
    #if [check_if_threaded] {set Environment threaded}
    if [check_if_pbs] {set Environment cluster}
    switch $Environment {
	cluster { set Njobs 48}
	threaded { set Njobs [check_num_cpu]}
	single {set Njobs 1}
    }
    set Charmmbin [lindex [check_cmd {imm1pb charmmpore charmm} ] 0]
    set Tclbin [lindex [check_cmd {tclsh8.6 tclsh}] 0]
    set Vmdbin [lindex [check_cmd {vmd}] 0]
    if ![info exists PdbPath] { set PdbPath [pwd]/pdb }
    if ![info exists JobPath] { set JobPath [pwd]/job }
    if ![info exists TemplatePath] {set TemplatePath [pwd]/template }
}

proc SimTask::printEnv {} {
    variable Njobs 
    variable Environment 
    variable Charmmbin 
    variable Vmdbin
    variable Tclbin
    variable PdbPath
    variable JobPath
    variable TemplatePath
    variable RunPath
    #setup 
    puts "Environment: $Environment"
    puts "Charmm Binary: $Charmmbin"
    puts "Tcl Binary: $Tclbin"
    puts "Vmd Binary: $Vmdbin"
    puts "Max number of concurrent jobs: $Njobs"
    puts "Place to look for PDB files: $PdbPath"
    puts "Place to look for template files: $TemplatePath"
    puts "Place to save job files: $JobPath"
    puts "Place to save data files: $RunPath"
}


proc SimTask::prepfiles { pdb_id files} {
 variable PdbPath
 variable RunPath
  if ![file exists $RunPath/$pdb_id ] { Path::mk_path $RunPath/$pdb_id }
 foreach file $files {
  if [file exists $PdbPath/$pdb_id/$file] {
   exec cp $PdbPath/$pdb_id/$file $RunPath/$pdb_id/$file
   } else {
    error "$file does not exist under $PdbPath/$pdb_id"
   }
 }
}

proc SimTask::parsefile { template options args } {
    upvar $options Options
    set extension [file extension $template]
    # default parsing pattern for charmm(.inp) or tcl(.tcl)
    switch $extension {
	.inp  {set set_variable_cmd "set \$var \$value"}
	default  {set set_variable_cmd "set \$var \$value"}
    }
    #drop seveal options for parsing
    set idx [lsearch $args --drop ]
    if { $idx >= 0} { set drop [lrange $args [expr $idx+1] end]} else {set drop {}}

    
    foreach var [Path::getVars Options] {
	if {[lsearch $var $drop] >=0 } {continue }
	upvar 1 $var value
	if ![info exists value] { set value [Path::getValue Options $var]}
	append inputfile [subst $set_variable_cmd]\n
    }
    set tfile [open $template r]
    append inputfile [read $tfile]
    close $tfile
    return $inputfile   
}

proc SimTask::random { run } {
    set r1 rand()
    set r2 rand()
    set r3 rand()
    return [expr int($r3*$r1*(10^int(100000*$r2)))]
}

proc SimTask::run { n base template options args } {
    variable RunPath
    variable TemplatePath
    variable Jobs
    variable JobPath
    upvar $options Options
    if { [string compare [file pathtype $base]  relative]==0 } { set base $RunPath/$base}
    set oldpath [pwd] 
    # add pre or post file system command 
    set pre_cmds ""
    set post_cmds ""
    set skip_if ""
    set combinejobby run
    set jobname [lindex [split $base /] end]
    set autopathpattern "array(key)_value/"
    set idx [lsearch $args --autopath ]
    if {$idx >=0 } {set autopathpattern [lindex $args [expr $idx +1]]}
    set idx [lsearch $args --pre ]
    if {$idx >=0 } {set pre_cmds [lindex $args [expr $idx +1]]}
    set idx [lsearch $args --post ]
    if {$idx >=0 } { set post_cmds [lindex $args [expr $idx+1]]}
    set idx [lsearch $args --skip ]
    if { $idx >=0 } { set skip_if [lindex $args [expr $idx+1]]}
    set idx [lsearch $args --combinejob ]
    if { $idx >=0 } {
	set idx2 [lsearch -start [expr 1+$idx] $args --*] 
	if {$idx2 <0 } {set combinejobby [lrange $args [expr $idx+1] end]} else {
	    set combinejobby [lrange $args [expr $idx+1] [expr $idx2-1]]}
    }
    set idx [lsearch $args --jobname ]
    if {$idx >=0} {set jobname [lindex $args [expr $idx+1]]}
    # add two more options to account n replicate runs
    set run {}
    set seed  { [ random @run ] }  
    for { set i 1 } { $i <=$n} { incr i} {lappend run run$i}
    Path::addVar Options run 
    Path::addVar Options seed
    #search for template file 
    set templatefilepath [Path::searchfiles $TemplatePath $template]
    if {[string length $templatefilepath] <=0 } { error "cannot find $template"}
    set templatebasename [lindex [split $template /] end]
    set template_extension [file extension $templatebasename]
    #set templatebasename [file rootname $templatebasename]

    #expand all variables 
    set jobID 0
    set jobs [dict create ]
    expand Options {
	# pattern use autopath ?
	set path [Path::autopath Options $autopathpattern --drop run seed]
	#if number of run =1, do not seperate into folder like run1, run2,...
	if {$n > 1} { set taskpath $base/$path/$run} else { set taskpath $base/$path }

	#prepare to add the simulation
	Path::mk_path $taskpath
	cd $taskpath	
	set passed 0
	#check if satisfy the --skip option, if true then do not do anything
	if ![catch {eval $skip_if } error ] { 
	    set passed [eval $skip_if ]
            #puts $taskpath
	    if { [string length $passed] <=0 } { set passed 0}
	} else {
	    puts "Warning error observed as:\n $error"
	    set passed 0
	}
	cd $base
	# if condition skip_if not passed or not defined, add the job
	if {!$passed || [string length $skip_if] == 0}  {
            # copy files
            foreach f [glob -type f $base/*] {
             puts "Copy $f into $taskpath"
             exec cp $f $taskpath/[lindex [split $f /] end]
            }
	    # make new input file 
	    puts "Create $template_extension file: $taskpath/$templatebasename"
	    set ofile [Path::open_file $taskpath/$templatebasename ]
	    puts $ofile [SimTask::parsefile $templatefilepath Options --drop run]
	    close $ofile	
	    
	    #add job to job queue
	    set task "cd $taskpath\n[subst $pre_cmds]\n[execuate $taskpath/$templatebasename]\n[subst $post_cmds]\ncd $RunPath\n"
	    # prepare the simulation job name
	    foreach var [Path::getVars Options] {
		set def [Path::getDef Options $var] 
		set value [Path::getValue Options $var]
		if {[string compare $def $value]!=0 && ![Path::isDependentVar $def Options]} {
		    dict set jobs $jobID $var $value
		}
	    }
	    dict set jobs $jobID path $taskpath
	    dict set jobs $jobID task $task
	    incr jobID
	}

    }
    #make job queue
    #puts $jobs
    if { [llength $jobs] <= 0 }  return 
    set queue [createJob $jobs run_$jobname $combinejobby]
    #execuate queue
    foreach job $queue {
	submit_or_wait $job
    }
    cd $oldpath
}

proc SimTask::ana { base dcdfile template option pattern args } { 
    variable TemplatePath
    variable RunPath
    upvar $option Options
    if { [string compare [file pathtype $base] relative]==0 } { set base $RunPath/$base}
    set oldpath [pwd]
    # if check the value contained in pre-defined range
    set checkrange 0
    if {[lsearch $args --checkrange]>=0} {set checkrange 1 }
    # add pre or post file system command 
    set pre_cmds ""
    set post_cmds ""
    set skip_if ""
    set combinejobby run
    set jobname [lindex [split $base /] end]
    #process command options    
    set idx [lsearch $args --pre ]
    if {$idx >=0 } {set pre_cmds [lindex $args [expr $idx +1]]}
    set idx [lsearch $args --post ]
    if {$idx >=0 } { set post_cmds [lindex $args [expr $idx+1]]}
    set idx [lsearch $args --skip ]
    if { $idx >=0 } { set skip_if [lindex $args [expr $idx+1]]}
    set idx [lsearch $args --combinejob ]
    if { $idx >=0 } {
	set idx2 [lsearch -start [expr 1+$idx] $args --*] 
	if {$idx2 <0 } {set combinejobby [lrange $args [expr $idx+1] end]} else {
	    set combinejobby [lrange $args [expr $idx+1] $idx2]}
    }
    set idx [lsearch $args --jobname ]
    if {$idx >=0} {set jobname [lindex $args [expr $idx+1]]}

    #search for template file 
    set templatefilepath [Path::searchfiles $TemplatePath $template]
    if {[string length $templatefilepath] <=0 } { error "cannot find $template"}
    set templatebasename [lindex [split $template /] end]
    set template_extension [file extension $templatebasename]

    #add fname option 
    set fname [file rootname [lindex [split $dcdfile /] end]]
    Path::addVar Options fname 
    
    #search the coordinate file need to be analyzed & generate processing job
    set jobID 0
    set jobs [dict create]
    foreach Path [Path::searchfiles $base $dcdfile] {
	set dir [file dirname $Path]
        set taskpath $dir 
	set fname [file rootname [lindex [split $Path /] end]]
        if [ catch {set vars [Path::parse $Path $pattern] } error ] { puts $error; continue}
	set doana 1
	puts "Parse $pattern from $dir: "
        #parse the variable based on pattern from dirname  
	foreach var $vars {
	    upvar 0 $var Var
	    upvar 1 $var OrigVar
            Path::setValue Options $var $Var
	    if {$checkrange } {
              puts "$var=$Var within? $OrigVar"
              if {[lsearch $OrigVar $Var] < 0}  { set doana 0; puts "Out of range, skip analysis" }
            }
	}
        #parse the rest of the variables based on the known variable
       foreach var [getVars Options] {
          if [Path::isDependentVar [Path::getDef Options $var] Options ] {
             Path::setValue Options $var [Path::evaluate Options $var [Path::getDef Options $var]]
             upvar 0 $var Var
             set Var [Path::getValue Options $var]
             continue
          }
          if [Path::isExpression [Path::getDef Options $var] ] {
             Path::setValue Options $var [Path::evaluate Options $var [Path::getValue Options $var]]
             upvar 0 $var Var
             set Var [Path::getValue Options $var]
             continue
          }
        }
        set passed 0
        cd $dir 
        #check if satisfy the --skip option, if true then do not do anything
        if ![catch {eval "expr $skip_if" } error ] {
            set passed [eval "expr $skip_if" ]
            if { [string length $passed] <=0 } { set passed 0}
            if { $passed} { puts "$skip_if satisfied. Analysis skipped."}
        } else {
            puts "Warning error observed as:\n $error"
            set passed 0
        }
        cd $base
        # if condition skip_if not passed or not defined, add the job
        # if no --skip option or specified --skip option and value in dirname is inside the range defined in "set varname { value1 value2 value3 }
	if {$doana && [expr !$passed || [string length $skip_if] == 0]} {
            # copy files necessary for simulation 
            #foreach f [glob -type f $base/*] {
            # puts "Copy $f into $taskpath"
            # exec cp $f $taskpath/[lindex [split $f /] end]
            #}
	    # create inp files for analysis
	    puts "Create $template_extension file: $dir/$templatebasename"
	    set ofile [Path::open_file $dir/$templatebasename ]
	    puts $ofile [SimTask::parsefile $templatefilepath Options]         
	    close $ofile 
	    #add job 
	    set task "cd $dir\n[subst $pre_cmds]\n[execuate $dir/$templatebasename]\n[subst $post_cmds]\ncd $RunPath\n"
	    foreach var $vars {
		upvar 0 $var Value
		dict set jobs $jobID $var $Value
	    }
	    dict set jobs $jobID path $dir
	    dict set jobs $jobID task $task
	    incr jobID
	}
    }
    #make job queue
    if { ![info exists jobs] } { return 0 }
    if { [llength $jobs] <= 0 } { return  0 }
    set queue [createJob $jobs ana_$jobname $combinejobby]
    #execuate queue
    foreach job $queue {
	submit_or_wait $job
    }
    cd $oldpath
}
proc SimTask::execuate { input } {
    variable RunPath
    variable Charmmbin
    set dir [file dirname $input]
    set fname [lindex [split $input /] end]
    set rootname [file rootname $fname]
    set extension [file extension $input]
    switch $extension {
	.inp { append cmd "$Charmmbin < $fname > $rootname.out \n" }
    }
    append cmd "grep -A 1 -niE 'terminating|warning|error' $rootname.out >&2 "
    return $cmd
}

proc SimTask::getDictValues { dict key } {
    set result {}
    if {[lsearch  [dict keys [dict get $dict 0]] $key] <0 } { return {}}
    foreach { id record} $dict {
	set value [dict get $record $key]
	if {[lsearch $result $value] <0} { lappend result $value }
    }
    return $result
}

proc SimTask::getDictID {dict args } {
    if {[llength $args] %2 == 0} {foreach {key value} $args { dict set toGet $key $value} } else {error "incorrect args! must be key & value pairs"}
    set result {}
    set keys {}
    set dictkeys [dict keys [dict get $dict 0]]
    foreach key [dict keys $toGet] {
	if {[lsearch $dictkeys $key] >=0} {lappend keys $key}
    }
    foreach { id record} $dict {
	set equalto 1
	foreach key $keys {
	    set value [dict get $record $key]
	    set value2 [dict get $toGet $key]
	    if {![Data::equal $value $value2]} {set equalto 0}
	}
	if {$equalto} {lappend result $id}
    }
    
    return $result
}
proc SimTask::getDictIDExcept { dict args } {
    set result {}
    set excludekeys {}
    if {[dict size $dict]<=0} { return {} }
    set dictkeys [dict keys [dict get $dict 0]]
    #puts $args
    foreach key $dictkeys {
	if {[lsearch $args $key] <0} {lappend excludekeys $key}
    }
    #puts $keys
    foreach { id record} $dict {
	set newkey ""
	foreach key $excludekeys {
	    set value [dict get $record $key]
	    append newkey "$key=$value,"
	}
	set newkey [string range $newkey 0 end-1]
	if [info exist V($newkey)] { lappend V($newkey) $id } else { set V($newkey) [list $id]}
    }
    return [array get V]
}
proc Data::equal { value value2} {
    if { [Data::isnumber $value] && [Data::isnumbe $value2]} {
	if {$value == $value2 } { return 1}
    } else {
	if {[string compare $value $value2]==0} { return 1}
    }
    return 0
}
set SimTask::JobHeader "#!/bin/bash\n#PBS -q production\n#PBS -N \$jobname\n#PBS -l select=1:ncpus=1\n#PBS -l place=pack\n#PBS -V"

proc SimTask::createJob { jobs jobname_prex combinejobby} {
    variable JobPath
    variable Environment
    variable JobHeader
    set today [clock format [clock seconds] -format %Y_%m_%d]
    set combinejobby [join $combinejobby " "]
    set queue {}
    foreach {name jobids} [eval "getDictIDExcept \{$jobs\}  $combinejobby path task seed" ] {
	set jobfilename  $JobPath/$today/${jobname_prex}_$name.job
	set jfile [Path::open_file $jobfilename]
	switch $Environment {
	    Cluster {
                set jobname ${jobname_prex}_$name
                set jobname [string range [string map { _ ""} $jobname] 0 14]
		puts $jfile [subst $JobHeader]
		foreach id $jobids {
		    puts $jfile [dict get [dict get $jobs $id] task]
		}

	    }
	    default {
                set jobname ${jobname_prex}_$name
                set jobname [string range [string map { _ ""} $jobname] 0 14]
		puts $jfile [subst $JobHeader]
		foreach id $jobids {
		    puts $jfile [dict get [dict get $jobs $id] task]
		}
	    }
	}
        close $jfile
	lappend queue $jobfilename
    }
    return $queue
}


SimTask::setup
